#include "framework.hh"
#include "entity.hh"
#include "flatbuffers_helper.hh"
#include "entity_search_cache.hh"
#include "../protocol/protocol_header.hh"
#include "../protocol/locator/register_entity_req_generated.h"
#include "../protocol/locator/register_entity_ack_generated.h"
#include "../protocol/locator/unregister_entity_req_generated.h"
#include "../protocol/locator/entity_relay_msg_req_generated.h"
#include "../protocol/locator/entity_search_req_generated.h"
#include "../protocol/locator/rpc_entity_req_generated.h"
#include "rpc_base.hh"

Framework::Framework()
    : sendBuffer_(new char[kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE]) {
}

Framework::~Framework() {
}

TimerHandle Framework::startTimer(std::time_t duration, std::uint64_t userData) {
    return ModuleRef.startTimer(duration, userData);
}

TimerHandle Framework::startTimerOnce(std::time_t duration, std::uint64_t userData) {
    return ModuleRef.startTimerOnce(duration, userData);
}

TimerHandle Framework::startTimer(std::time_t duration, TimerCallback cb, std::uint64_t userData) {
    return ModuleRef.startTimer(duration, cb, userData);
}

TimerHandle Framework::startTimerOnce(std::time_t duration, TimerCallback cb, std::uint64_t userData) {
    return ModuleRef.startTimerOnce(duration, cb, userData);
}

void Framework::cancelTimer(TimerHandle timerID) {
    ModuleRef.cancelTimer(timerID);
}

bool Framework::findPath(SiteType siteType, UserData userData) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    return ModuleRef.getEngine()->findPath(siteType, userData);
}

void Framework::closePath(PathID pathID) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return;
    }
    ModuleRef.getEngine()->closePath(pathID);
}

void Framework::closeSession(SessionID sessionID) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return;
    }
    ModuleRef.getEngine()->closeSession(sessionID);
}

bool Framework::talkPath(PathID pathID, const Message& msg) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    auto msgID = ProtocolRegisterRef.getMessageID(msg);
    FRAMEWORK_CHECK_RETURN(!msgID, false, "Message not found, need register first:" << msg.GetDescriptor()->name());
    auto bytes = msg.ByteSize();
    auto packetSize = bytes + sizeof(ProtocolHeader);
    FRAMEWORK_CHECK_RETURN(packetSize > kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE,
        false, "Packet size invalid, msgID:" << msgID << ", size:" << bytes);
    buildHeaderInSendBuffer(0, msgID, packetSize);
    FRAMEWORK_CHECK_RETURN(!msg.SerializeToArray(sendBuffer_.get() + sizeof(ProtocolHeader), bytes),
        false, "SerializeToArray failed, msgID:" << msgID << ", size:" << bytes);
    return ModuleRef.getEngine()->talkPath(pathID, sendBuffer_.get(), packetSize);
}

bool Framework::talkPath(PathID pathID, const char * data, std::size_t size) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    return ModuleRef.getEngine()->talkPath(pathID, data, size);
}

bool Framework::talkSession(SessionID sessionID, const Message& msg) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    auto msgID = ProtocolRegisterRef.getMessageID(msg);
    FRAMEWORK_CHECK_RETURN(!msgID, false, "Message not found, need register first:" << msg.GetDescriptor()->name());
    auto bytes = msg.ByteSize();
    auto packetSize = bytes + sizeof(ProtocolHeader);
    FRAMEWORK_CHECK_RETURN(packetSize > kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE,
        false, "Packet size invalid, msgID:" << msgID << ", size:" << bytes)
        buildHeaderInSendBuffer(0, msgID, packetSize);
    FRAMEWORK_CHECK_RETURN(!msg.SerializeToArray(sendBuffer_.get() + sizeof(ProtocolHeader), bytes),
        false, "SerializeToArray failed, msgID:" << msgID << ", size:" << bytes);
    return ModuleRef.getEngine()->talkSession(sessionID, sendBuffer_.get(), packetSize);
}

bool Framework::talkSession(SessionID sessionID, const char * data, std::size_t size) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    return ModuleRef.getEngine()->talkSession(sessionID, data, size);
}

bool Framework::talkSite(const kratos::grid::SiteID & siteID, const Message & msg) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    auto msgID = ProtocolRegisterRef.getMessageID(msg);
    FRAMEWORK_CHECK_RETURN(!msgID, false, "Message not found, need register first:" << msg.GetDescriptor()->name());
    auto bytes = msg.ByteSize();
    auto packetSize = bytes + sizeof(ProtocolHeader);
    FRAMEWORK_CHECK_RETURN(packetSize > kratos::NETWORK_IOLOOP_RECV_BUFFER_MAX_SIZE, false,
        "Packet size invalid, msgID:" << msgID << ", size:" << bytes);
    buildHeaderInSendBuffer(0, msgID, packetSize);
    FRAMEWORK_CHECK_RETURN(!msg.SerializeToArray(sendBuffer_.get() + sizeof(ProtocolHeader), bytes),
        false, "SerializeToArray failed, msgID:" << msgID << ", size:" << bytes);
    return ModuleRef.getEngine()->talkSite(siteID, sendBuffer_.get(), packetSize);
}

bool Framework::talkSite(const kratos::grid::SiteID & siteID, const char * data, std::size_t size) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    return ModuleRef.getEngine()->talkSite(siteID, data, size);
}

bool Framework::talkEntity(std::uint64_t from, std::uint64_t to, const Message& msg) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    auto msgID = ProtocolRegisterRef.getMessageID(msg);
    FRAMEWORK_CHECK_RETURN(!msgID, false, "Message not found, need register first:" << msg.GetDescriptor()->name());
    auto entity = EntityManagerRef.getEntity(to);
    if (entity) {
        // For the local entity
        auto handler = ProtocolRegisterRef.getHandler(msgID);
        if (handler) {
            handler->onEntityMessage(from, to, const_cast<Message&>(msg));
            return true;
        } else {
            return EntityManagerRef.onMessage(from, to, msg);
        }
    } else {
        // For the remote entity
        auto pathID = ModuleRef.getEngine()->getLocator(to);
        FRAMEWORK_CHECK_RETURN(INVALID_PATH_ID == pathID, false,
            "Entity not found in locator, entity ID:" << to);
        BUILD_FLATBUFFER(EntityRelayMsgReq, fb, from, to, msg.ByteSize());
        return talkLocator(pathID, getMsgRealID(LocatorMsgID::RELAY_REQ), fb, msgID, msg);
    }
}

void Framework::broadcastEntity(std::uint32_t entityType, const Message& msg, std::uint64_t exceptID) {
    auto& entityIDs = EntitySearchCacheManagerRef.getEntityIDs(entityType);
    for (auto id : entityIDs) {
        if (id != exceptID) {
            talkEntity(id, msg);
        }
    }
}

bool Framework::talkEntity(std::uint64_t to, const Message & msg) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    return talkEntity(0, to, msg);
}

bool Framework::registerEntity(std::uint64_t entityID, std::uint32_t entityType) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    FRAMEWORK_CHECK_RETURN(!tryToStartRegisterTimer(entityID), false,
        "Start registering timer failed, entity ID:" << entityID << ", entity Type:" << entityType);
    pendingEntityMap_.emplace(entityID, EntityRegistry{ entityID, entityType, INVALID_PATH_ID, 0 });
    auto pathID = ModuleRef.getEngine()->getLocator(entityID);
    FRAMEWORK_CHECK_RETURN(pathID == INVALID_PATH_ID, false,
        "The locator is not ready for registering for entity ID:" << entityID << ", entity Type:" << entityType);
    BUILD_FLATBUFFER(RegisterEntityReq, fb, entityID, entityType);
    return talkLocator(pathID, getMsgRealID(LocatorMsgID::REGISTER_REQ), fb);
}

bool Framework::unregisterEntity(std::uint64_t entityID) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    auto pendingIt = pendingEntityMap_.find(entityID);
    if (pendingIt != pendingEntityMap_.end()) {
        // not finish registering
        pendingEntityMap_.erase(pendingIt);
        return true;
    }
    PathID pathID = INVALID_PATH_ID;
    auto registeredIt = registeredEntityMap_.find(entityID);
    if (registeredIt != registeredEntityMap_.end()) {
        // Remove from local and locator
        pathID = registeredIt->second.pathID;
        FRAMEWORK_CHECK_RETURN(pathID == INVALID_PATH_ID, false, "Registered entity but owner's path ID invalid, entity ID:" << entityID);
        registeredEntityMap_.erase(registeredIt);
        BUILD_FLATBUFFER(UnregisterEntityReq, fb, entityID);
        return talkLocator(pathID, getMsgRealID(LocatorMsgID::UNREGISTER_REQ), fb);
    } else {
        errorlog << "Try to un-register entity but failed, entity ID:" << entityID << endlog;
        return false;
    }
}

bool Framework::searchEntity(std::uint32_t entityType, SearchMethod method, SearchType searchType, std::time_t timeout,
    std::uint64_t userData) {
    if (ModuleType::SERVICE_MODULE == onModuleType()) {
        return false;
    }
    return EntitySearchCacheManagerRef.search(entityType, searchType, method, timeout, userData);
}

khttp::HttpRequest * Framework::newHttpRequest(HttpMethod cb, const std::string & host, const std::string & uri, int32_t port) {
    return ModuleRef.newHttpRequest(cb, host, uri, port);
}

bool Framework::newHttpServer(HttpMethod cb, const std::string & host, int32_t port) {
    return ModuleRef.newHttpServer(cb, host, port);
}

std::string Framework::doCommand(const std::string & moduleName, const std::string & command) {
    return ModuleRef.getEngine()->doCommand(moduleName, command);
}

std::string Framework::doCommand(const std::string & command) {
    return ModuleRef.getEngine()->doCommand(command);
}

const std::string& Framework::getRootPath() {
    return ModuleRef.getRootPath();
}

bool Framework::connect(const std::string& host, std::uint64_t userKey, std::time_t timeout) {
    return ModuleRef.getEngine()->connect(host, userKey, timeout);
}
